import numpy as np
import faiss
from sklearn.metrics import accuracy_score

class FaissKNeighbors:
    def __init__(self, k=1, res=None):
        self.index = None  # 用于存储训练数据的索引
        self.y = None      # 用于存储训练数据的标签
        self.k = k         # 最近邻个数
        self.res = res     # FAISS GPU资源对象

    def fit(self, X, y):
        d = X.shape[1]
        # 初始化IndexFlatL2索引（欧氏距离度量）
        self.index = faiss.IndexFlatL2(d)
        # 若有GPU资源，将索引迁移到GPU
        if self.res is not None:
            self.index = faiss.index_cpu_to_gpu(self.res, 0, self.index)
        # 存储标签并添加训练数据到索引
        self.y = y
        self.index.add(X.astype(np.float32))

    def predict(self, X):
        # 搜索k个最近邻
        distances, indices = self.index.search(X.astype(np.float32), self.k)
        # 投票确定预测标签
        predictions = []
        for idx in indices:
            labels = self.y[idx]
            unique, counts = np.unique(labels, return_counts=True)
            predictions.append(unique[np.argmax(counts)])
        return np.array(predictions)

    def score(self, X, y_true):
        # 计算预测准确率
        y_pred = self.predict(X)
        return accuracy_score(y_true, y_pred)

# 测试代码
if __name__ == "__main__":
    # 生成模拟数据
    np.random.seed(42)
    X_train = np.random.rand(100, 10).astype(np.float32)
    y_train = np.random.randint(0, 2, 100)
    X_test = np.random.rand(20, 10).astype(np.float32)
    y_test = np.random.randint(0, 2, 20)

    # 初始化并训练模型
    knn = FaissKNeighbors(k=3)
    knn.fit(X_train, y_train)

    # 预测并计算准确率
    y_pred = knn.predict(X_test)
    accuracy = knn.score(X_test, y_test)
    print("预测结果:", y_pred)
    print("准确率:", accuracy)